#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pymysql
# from sshtunnel import SSHTunnelForwarder
import logging
import traceback
import unittest


# 似乎无法输出日志?
class DbMysql(object):
    def __init__(self, host, port, user, password, db):
        self.my_host = host
        self.my_port = port
        self.my_user = user
        self.my_password = password
        self.my_db = db
        self.my_conn = None
        self.ssh_host = None
        self.ssh_port = None
        self.ssh_user = None
        self.ssh_password = None
        self.server = None

    # def set_ssh_tunnel(self, host, port, user, password):
    #     self.ssh_host = host
    #     self.ssh_port = port
    #     self.ssh_user = user
    #     self.ssh_password = password
    #     self.server = SSHTunnelForwarder(
    #         (self.ssh_host, self.ssh_port),  # B机器的配置
    #         ssh_password=self.ssh_password,
    #         ssh_username=self.ssh_user,
    #         remote_bind_address=(self.my_host, self.my_port)
    #     )  # A机器的配置
    #     self.server.start()

    def connect(self):
        try:
            if self.server is not None:
                self.my_conn = pymysql.connect(
                    host='127.0.0.1',  # 此处必须是是127.0.0.1
                    port=self.server.local_bind_port,
                    user=self.my_user,
                    passwd=self.my_password,
                    db=self.my_db,
                    charset='utf8'
                )
            else:
                self.my_conn = pymysql.connect(
                    host=self.my_host,
                    user=self.my_user,
                    passwd=self.my_password,
                    port=self.my_port,
                    db=self.my_db,
                    charset="utf8"
                )
            logging.info('connect to msyql {0}'.format(self.my_conn))
        except Exception as err:
            logging.error('get_mdb_conn connect error:{0}'.format(err))
            self.my_conn = None
        return self.my_conn

    def execute(self, sqls):
        status_code = 0
        if not isinstance(sqls, list):
            sqls = [sqls]
        if self.my_conn is None:
            self.connect()
        with self.my_conn.cursor() as cursor:
            i = 0
            retry = False
            while i < len(sqls):
                sql = sqls[i]
                try:
                    cursor.execute(sql)
                    logging.info(sql)
                except Exception as e:
                    if e.args[0] == 1062:
                        # 重复继续执行
                        pass
                    elif e.args[0] in [2006, 2013] or e.args[0] == "(0, '')":
                        if not retry:
                            retry = True
                            self.connect()
                            cursor = self.my_conn.cursor()
                            # 重新连接 再次执行
                            continue
                        else:
                            status_code += -1
                    else:
                        logging.error(e)
                        logging.error(sql)
                        logging.error(traceback.format_exc())
                        status_code += -1
                i += 1
                if retry:
                    retry = False
        return status_code

    def query(self, sql):
        rows = []
        if self.my_conn is None:
            self.connect()
        with self.my_conn.cursor() as cursor:
            while True:
                try:
                    cursor.execute(sql)
                    rows = cursor.fetchall()
                except Exception as e:
                    if e.args[0] in [2006, 2013] or e.args[0] == "(0, '')":
                        self.connect()
                        cursor = self.my_conn.cursor()
                        # 重新连接 再次执行
                        continue
                    else:
                        logging.error(e)
                        logging.error(sql)
                        logging.error(traceback.format_exc())
                break
        return rows

    def commit(self):
        if self.my_conn is not None:
            self.my_conn.commit()

    def close(self):
        if self.my_conn is not None:
            self.my_conn.close()
        if self.server is not None:
            self.server.close()


class TestDbMysql(unittest.TestCase):
    def setUp(self) -> None:
        tz_host = '127.0.0.1'
        tz_port = 3306
        tz_user = 'jim'
        tz_password = '111111'
        tz_db = 'app'
        self.mdb = DbMysql(tz_host, tz_port, tz_user, tz_password, tz_db)

    def tearDown(self) -> None:
        self.mdb.close()

    def test_mdb(self):
        sql = '''show tables;'''
        rows = self.mdb.query(sql)
        for row in rows:
            print(row)
        self.assertNotEqual(len(rows), 0, 'no tables in db?')


if __name__ == '__main__':
    unittest.main()
